"""
初始数据加载模块
从data/init目录读取Excel文件并生成贝叶斯优化所需的初始数据集
"""

import logging
from pathlib import Path

import pandas as pd

from src.metric import analyze_datafile

# 设置日志
logging.basicConfig(
    level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)


def load_initial_data_from_directory(
    data_dir: str,
) -> list[tuple[dict, dict]]:
    """
    从指定目录加载初始数据并生成贝叶斯优化所需的格式

    Args:
        data_dir: 包含初始数据Excel文件的目录路径
    Returns:
        preexisting_trials: 包含全部参数和全部目标值的列表
            格式: [({"t1": 25, "t2": 25, ...}, {"efficiency": 1.74, ...}), ...]
    """
    data_path = Path(data_dir)

    # 检查目录是否存在
    if not data_path.exists():
        logger.error(f"目录不存在: {data_path}")
        raise FileNotFoundError(f"目录不存在: {data_path}")

    # 获取所有Excel文件
    excel_files = list(data_path.glob("*.xlsx")) + list(
        data_path.glob("*.xls")
    )

    if not excel_files:
        logger.error(f"在 {data_path} 中未找到Excel文件")
        raise FileNotFoundError(f"在 {data_path} 中未找到Excel文件")

    logger.info(f"找到 {len(excel_files)} 个Excel文件")

    # 存储数据
    preexisting_trials = []

    # 处理每个文件
    for excel_file in excel_files:
        try:
            logger.info(f"正在处理: {excel_file.name}")

            # 使用data_analysis模块的函数解析文件, 并添加到列表
            params, metrics = analyze_datafile(excel_file)
            preexisting_trials.append((params, metrics))

        except Exception as e:
            logger.error(f"处理文件 {excel_file.name} 时出错: {e}")
            continue

    # 检查是否有数据
    if not preexisting_trials:
        logger.error("没有成功处理任何文件")
        raise ValueError("没有成功处理任何文件")

    logger.info(f"\n成功加载 {len(preexisting_trials)} 个初始数据点")

    return preexisting_trials


def save_initial_data(
    preexisting_trials: list[tuple[dict, dict]], save_dir: str
) -> None:
    """
    将初始数据保存为Excel文件，包含两个工作表：
    - data: 实际数据
    - variable: 变量分类信息（params或metrics）

    Args:
        preexisting_trials: 包含参数和目标值的列表
        save_dir: 保存目录
    """
    save_path = Path(save_dir)
    save_path.mkdir(parents=True, exist_ok=True)

    # 将数据转换为DataFrame格式
    data_rows = []
    param_names = set()
    metric_names = set()

    for params, metrics in preexisting_trials:
        # 合并参数和指标到一行
        row = {**params, **metrics}
        data_rows.append(row)
        # 记录参数和指标的名称
        param_names.update(params.keys())
        metric_names.update(metrics.keys())

    # 创建数据DataFrame
    df_data = pd.DataFrame(data_rows)

    # 创建变量分类DataFrame
    variable_info = []
    for col in df_data.columns:
        if col in param_names:
            variable_info.append({"variable": col, "type": "params"})
        elif col in metric_names:
            variable_info.append({"variable": col, "type": "metrics"})

    df_variable = pd.DataFrame(variable_info)

    # 保存为Excel格式，包含两个工作表
    excel_file = save_path / "preexisting_trials.xlsx"
    with pd.ExcelWriter(excel_file, engine="openpyxl") as writer:
        df_data.to_excel(writer, sheet_name="data", index=False)
        df_variable.to_excel(writer, sheet_name="variable", index=False)

    logger.info("\n初始数据已保存:")
    logger.info(f"  文件路径: {excel_file}")
    logger.info(f"  数据点数: {len(preexisting_trials)}")


def load_saved_initial_data(
    load_dir: str,
) -> list[tuple[dict, dict]]:
    """
    从已保存的Excel文件加载初始数据
    从 data 工作表读取数据，从 variable 工作表读取变量分类信息

    Args:
        load_dir: 数据文件所在目录

    Returns:
        preexisting_trials: 包含参数和目标值的列表
    """

    load_path = Path(load_dir)
    excel_file = load_path / "preexisting_trials.xlsx"

    if not excel_file.exists():
        logger.error(f"初始数据文件不存在: {excel_file}")
        raise FileNotFoundError(f"初始数据文件不存在: {excel_file}")

    # 从data工作表读取数据
    df_data = pd.read_excel(excel_file, sheet_name="data", engine="openpyxl")

    # 从variable工作表读取变量分类信息
    df_variable = pd.read_excel(
        excel_file, sheet_name="variable", engine="openpyxl"
    )

    # 构建变量分类字典
    variable_types = {}
    for _, row in df_variable.iterrows():
        variable_types[row["variable"]] = row["type"]

    # 转换为正确格式
    preexisting_trials = []

    for _, row in df_data.iterrows():
        row_dict = row.to_dict()

        params = {}
        metrics = {}

        for key, value in row_dict.items():
            # 根据variable工作表的信息分类, 并将参数转换为整数
            var_type = variable_types.get(key, "unknown")
            if var_type == "params":
                params[key] = int(value)
            elif var_type == "metrics":
                metrics[key] = value

        preexisting_trials.append((params, metrics))

    logger.info("\n从文件加载初始数据:")
    logger.info(f"  文件路径: {excel_file}")
    logger.info(f"  数据点数: {len(preexisting_trials)}")

    return preexisting_trials


def load_preexisting_trial(
    data_dir: str,
    parameters: list[str],
    optimization_metric: str,
    motor_rpm: int,
    force_reload: bool = False,
) -> list[tuple[dict, dict]]:
    """
    准备初始数据（优先从已保存文件加载，否则重新处理）

    Args:
        data_dir: 数据目录
        parameters: 需要保留的参数列表，例如 ["t1", "t2", "t3", "t4"]
        optimization_metric: 优化目标指标
        motor_rpm: 电机转速
        force_reload: 是否强制重新处理数据

    Returns:
        preexisting_trials: 包含参数和目标值的列表
    """
    data_path = Path(data_dir)
    excel_file = data_path / "preexisting_trials.xlsx"

    # 如果已存在保存的数据且不强制重新加载
    if not force_reload and excel_file.exists():
        logger.info("从已保存的文件加载初始数据...")
        preexisting_trials = load_saved_initial_data(data_dir)
    else:
        # 重新处理数据
        logger.info("从Excel文件处理初始数据...")
        preexisting_trials = load_initial_data_from_directory(data_dir)
        # 保存处理后的数据
        save_initial_data(preexisting_trials, data_dir)

    # 过滤参数、指标和电机转速
    logger.info(f"过滤参数，只保留: {parameters}")
    logger.info(f"过滤指标，只保留: {optimization_metric}")
    logger.info(f"过滤电机转速，只保留: {motor_rpm}rpm")

    filtered_trials = [
        (
            {k: v for k, v in params.items() if k in parameters},
            {
                f"{optimization_metric}_{motor_rpm}rpm": metrics[
                    optimization_metric
                ]
            },
        )
        for params, metrics in preexisting_trials
        if params.get("rpm") == motor_rpm  # 只保留匹配转速的数据
    ]

    logger.info(f"过滤后剩余 {len(filtered_trials)} 个数据点")

    return filtered_trials
